#include <lio_ndt/front_end/front_end.hpp>
#include <cmath>
#include <pcl/common/transforms.h>
#include <glog/logging.h>
#include <boost/make_shared.hpp>

#include "lio_ndt/method/registration.h"
#include "lio_ndt/method/pcl_icp.h"
#include "lio_ndt/method/pcl_ndt.h"
#include "lio_ndt/method/optimized_ICP_GN.h"

namespace lio_ndt {

FrontEnd::FrontEnd() 
    : local_map_ptr_(boost::make_shared<CloudData::CLOUD>())
    , global_map_ptr_(boost::make_shared<CloudData::CLOUD>())
    , result_cloud_ptr_(boost::make_shared<CloudData::CLOUD>())
    , init_pose_(Eigen::Matrix4f::Identity())
    , has_new_local_map_(false)
    , has_new_global_map_(false) {
    FLAGS_logtostderr = 1;

    std::string method_type = "pcl_ndt";  
    if (method_type == "pcl_icp") {
        LOG(INFO) << "[FrontEnd] Selected Registration Method: PCL ICP";
        LOG(INFO) << "----------------------------------------------------------------------";
        
        matcher_ = boost::make_shared<PCLICP>();
        auto icp_ptr = boost::dynamic_pointer_cast<PCLICP>(matcher_);
        if (icp_ptr) {
            icp_ptr->SetMaxCorrespondenceDistance(1.0);
            icp_ptr->SetMaximumIterations(50);
            icp_ptr->SetTransformationEpsilon(1e-6);
        }
    } 
    else if (method_type == "pcl_ndt") {
        LOG(INFO) << "[FrontEnd] Selected Registration Method: PCL NDT";
        LOG(INFO) << "----------------------------------------------------------------------";
        
        matcher_ = boost::make_shared<PCLNDT>();
        auto ndt_ptr = boost::dynamic_pointer_cast<PCLNDT>(matcher_);
        if (ndt_ptr) {
            ndt_ptr->SetResolution(1.0);  // 降低分辨率，提高配准精度
            ndt_ptr->SetMaximumIterations(50);
            ndt_ptr->SetTransformationEpsilon(1e-3);
        }
    } 
    else {
        LOG(INFO) << "[FrontEnd] Selected Registration Method: Optimized GN-ICP";
        LOG(INFO) << "----------------------------------------------------------------------";
        
        opti_icp_.reset(new OptimizedICPGN());
        opti_icp_->SetMaxCorrespondDistance(1.0);
        opti_icp_->SetMaxIterations(30);
        opti_icp_->SetTransformationEpsilon(0.01);
    }

    // 减小体素滤波分辨率，保留更多点云细节
    cloud_filter_.setLeafSize(1.0f, 1.0f, 1.0f);
    local_map_filter_.setLeafSize(0.5f, 0.5f, 0.5f);
    display_filter_.setLeafSize(0.5f, 0.5f, 0.5f);
}

Eigen::Matrix4f FrontEnd::Update(const CloudData &cloud_data) {
    current_frame_.cloud_data.time = cloud_data.time;
    current_frame_.cloud_data.cloud_ptr = boost::make_shared<CloudData::CLOUD>(*cloud_data.cloud_ptr);

    std::vector<int> indices;
    pcl::removeNaNFromPointCloud(*current_frame_.cloud_data.cloud_ptr, *current_frame_.cloud_data.cloud_ptr, indices);

    CloudData::CLOUD_PTR filtered_cloud_ptr = boost::make_shared<CloudData::CLOUD>();
    cloud_filter_.setInputCloud(current_frame_.cloud_data.cloud_ptr);
    cloud_filter_.filter(*filtered_cloud_ptr);

    static Eigen::Matrix4f step_pose = Eigen::Matrix4f::Identity();
    static Eigen::Matrix4f last_pose = init_pose_;
    static Eigen::Matrix4f predict_pose = init_pose_;
    static Eigen::Matrix4f last_key_frame_pose = init_pose_;

    if (local_map_frames_.empty()) {
        current_frame_.pose = init_pose_;
        UpdateNewFrame(current_frame_);
        return current_frame_.pose;
    }

    bool success = false;
    if (matcher_) {
        success = matcher_->ScanMatch(filtered_cloud_ptr, predict_pose, result_cloud_ptr_, current_frame_.pose);
    } else if (opti_icp_) {
        success = opti_icp_->Match(filtered_cloud_ptr, predict_pose, result_cloud_ptr_, current_frame_.pose);
    }

    if (!success) {
        LOG(WARNING) << "[FrontEnd] ScanMatch failed, using predict pose";
        current_frame_.pose = predict_pose;
    }

    std::cout << "fitness score: ";
    if (matcher_) {
        double score = matcher_->GetFitnessScore();
        if (score < 0) {
            std::cout << "(no score available yet)" << std::endl;
        } else if (!std::isfinite(score)) {
            std::cout << "(not converged / inf)" << std::endl;
        } else {
            std::cout << score << std::endl;
        }
    } else if (opti_icp_) {
        std::cout << opti_icp_->GetFitnessScore() << std::endl;
    } else {
        std::cout << "(no matcher)" << std::endl;
    }

    step_pose = last_pose.inverse() * current_frame_.pose;
    predict_pose = current_frame_.pose * step_pose;
    last_pose = current_frame_.pose;

    // 减小关键帧移动距离阈值，更频繁生成关键帧
    float move_distance = fabs(last_key_frame_pose(0, 3) - current_frame_.pose(0, 3)) +
                          fabs(last_key_frame_pose(1, 3) - current_frame_.pose(1, 3)) +
                          fabs(last_key_frame_pose(2, 3) - current_frame_.pose(2, 3));
    if (move_distance > 1.0) {  // 从2.0米调整为1.0米
        UpdateNewFrame(current_frame_);
        last_key_frame_pose = current_frame_.pose;
    }

    return current_frame_.pose;
}

bool FrontEnd::SetInitPose(const Eigen::Matrix4f &init_pose) {
    init_pose_ = init_pose;
    return true;
}

void FrontEnd::UpdateNewFrame(const Frame &new_key_frame) {
    Frame key_frame = new_key_frame;
    key_frame.cloud_data.cloud_ptr = boost::make_shared<CloudData::CLOUD>(*new_key_frame.cloud_data.cloud_ptr);

    CloudData::CLOUD_PTR transformed_cloud_ptr = boost::make_shared<CloudData::CLOUD>();

    local_map_frames_.push_back(key_frame);
    while (local_map_frames_.size() > 20) {
        local_map_frames_.pop_front();
    }

    local_map_ptr_ = boost::make_shared<CloudData::CLOUD>();
    for (auto &frame : local_map_frames_) {
        pcl::transformPointCloud(*frame.cloud_data.cloud_ptr, *transformed_cloud_ptr, frame.pose);
        *local_map_ptr_ += *transformed_cloud_ptr;
    }
    has_new_local_map_ = true;

    CloudData::CLOUD_PTR target_map = boost::make_shared<CloudData::CLOUD>();
    if (local_map_frames_.size() < 10) {
        target_map = local_map_ptr_;
    } else {
        local_map_filter_.setInputCloud(local_map_ptr_);
        local_map_filter_.filter(*target_map);
    }

    if (matcher_) matcher_->SetTargetCloud(target_map);
    if (opti_icp_) opti_icp_->SetTargetCloud(target_map);

    global_map_frames_.push_back(key_frame);
    // 减小全局地图更新间隔，更频繁生成全局地图
    const size_t kGlobalUpdateInterval = 20;  // 从100帧调整为20帧
    if (global_map_frames_.size() % kGlobalUpdateInterval == 0) {
        global_map_ptr_ = boost::make_shared<CloudData::CLOUD>();
        for (auto &frame : global_map_frames_) {
            pcl::transformPointCloud(*frame.cloud_data.cloud_ptr, *transformed_cloud_ptr, frame.pose);
            *global_map_ptr_ += *transformed_cloud_ptr;
        }
        has_new_global_map_ = true;
    }
}

bool FrontEnd::GetNewLocalMap(CloudData::CLOUD_PTR &local_map_ptr) {
    if (has_new_local_map_) {
        display_filter_.setInputCloud(local_map_ptr_);
        display_filter_.filter(*local_map_ptr);
        has_new_local_map_ = false;
        return true;
    }
    return false;
}

bool FrontEnd::GetNewGlobalMap(CloudData::CLOUD_PTR &global_map_ptr) {
    if (has_new_global_map_) {
        display_filter_.setInputCloud(global_map_ptr_);
        display_filter_.filter(*global_map_ptr);
        has_new_global_map_ = false;
        return true;
    }
    return false;
}

bool FrontEnd::GetCurrentScan(CloudData::CLOUD_PTR &current_scan_ptr) {
    display_filter_.setInputCloud(result_cloud_ptr_);
    display_filter_.filter(*current_scan_ptr);
    return true;
}

}  // namespace lio_ndt
